#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time
import copy

sys.path.append('.')

"""
    将相同类型图片合并的可视化工具，一次显示一个设备的所有图片，选择相同类型的图片合并（主要是相同的人）
"""

video_base_paths = [
    # r'/rootfs/media/kasim/DataSet/RecordEx',
    # r'/rootfs/media/kasim/DataSet/Record',
    # r'/rootfs/media/kasim/DataSet/Record2',
    r'/rootfs/media/yery/Dataset/RecordEx',
    r'/rootfs/media/yery/Dataset/Record',
    r'/rootfs/media/yery/Dataset/Record2',
]

DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

h_image_count = 12
v_image_count = 10
img_width = 200
img_height = 200

batch_size = 32


ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()


def get_xy_bbox_idx(x, y):
    global image_bboxes
    bbox_idx = None
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def IOU(bbox, bboxes):
    bboxes2 = np.array(bboxes, dtype=np.int32)
    bboxes1 = np.array([bbox], dtype=np.int32)
    # bboxes1 = np.array([bbox], dtype=np.int32).repeat(bboxes2.shape[0], axis=0)
    lt = np.maximum(bboxes1[:, :2], bboxes2[:, :2])
    rb = np.minimum(bboxes1[:, 2:], bboxes2[:, 2:])
    wh = (rb - lt)
    wh = np.maximum(wh, 0)
    overlap = wh[:, 0] * wh[:, 1]
    area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1])
    area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (bboxes2[:, 3] - bboxes2[:, 1])
    ious = overlap / (area1 + area2 - overlap)
    return ious


def get_xy_bbox_idxes(x1, y1, x2, y2):
    global image_bboxes
    if x1 < x2:
        sx = x1
        ex = x2
    else:
        sx = x2
        ex = x1
    if y1 < y2:
        sy = y1
        ey = y2
    else:
        sy = y2
        ey = y1
    iou = IOU([sx, sy, ex, ey], image_bboxes)
    bbox_idxes = np.nonzero(iou > 0)[0].tolist()
    return bbox_idxes


def draw_rect(image, start_xy, end_xy):
    if (start_xy is None) or (end_xy is None):
        return
    x1, y1 = start_xy
    x2, y2 = end_xy
    if x1 < x2:
        sx = x1
        ex = x2
    else:
        sx = x2
        ex = x1
    if y1 < y2:
        sy = y1
        ey = y2
    else:
        sy = y2
        ey = y1
    cv2.rectangle(image, (sx, sy), (ex, ey), (0, 255, 0), thickness=2)


start_bbox_init = False
start_modify_bbox = False
start_mouse_xy = None
end_mouse_xy = None
start_bbox_init2 = False
start_modify_bbox2 = False
start_mouse_xy2 = None
end_mouse_xy2 = None
current_mouse_xy = None


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes, start_bbox_init, start_modify_bbox, start_mouse_xy, \
        end_mouse_xy, start_bbox_init2, start_modify_bbox2, start_mouse_xy2, end_mouse_xy2, current_mouse_xy

    if draw_image is not None:
        current_mouse_xy = [x, y]
        if event == cv2.EVENT_MOUSEMOVE and flags == (cv2.EVENT_FLAG_SHIFTKEY + cv2.EVENT_FLAG_LBUTTON):
            if start_modify_bbox:
                end_mouse_xy = [x, y]
                bbox_idxes = get_xy_bbox_idxes(start_mouse_xy[0], start_mouse_xy[1], x, y)
                for bbox_idx in bbox_idxes:
                    if bbox_idx not in select_image_bbox_idxes:
                        select_image_bbox_idxes.add(bbox_idx)
            return

        if event == cv2.EVENT_LBUTTONDOWN and flags == (cv2.EVENT_FLAG_SHIFTKEY + cv2.EVENT_FLAG_LBUTTON):
            if not start_bbox_init:
                start_bbox_init = True
                start_modify_bbox = True
                start_mouse_xy = [x, y]
            return

        start_bbox_init = False
        start_modify_bbox = False
        start_mouse_xy = None
        end_mouse_xy = None

        if event == cv2.EVENT_MOUSEMOVE and flags == (cv2.EVENT_FLAG_CTRLKEY + cv2.EVENT_FLAG_LBUTTON):
            if start_modify_bbox2:
                end_mouse_xy2 = [x, y]
                bbox_idxes = get_xy_bbox_idxes(start_mouse_xy2[0], start_mouse_xy2[1], x, y)
                for bbox_idx in bbox_idxes:
                    if bbox_idx in select_image_bbox_idxes:
                        select_image_bbox_idxes.remove(bbox_idx)
            return

        if event == cv2.EVENT_LBUTTONDOWN and flags == (cv2.EVENT_FLAG_CTRLKEY + cv2.EVENT_FLAG_LBUTTON):
            if not start_bbox_init2:
                start_bbox_init2 = True
                start_modify_bbox2 = True
                start_mouse_xy2 = [x, y]
            return

        start_bbox_init2 = False
        start_modify_bbox2 = False
        start_mouse_xy2 = None
        end_mouse_xy2 = None

        if event == cv2.EVENT_LBUTTONDOWN:
            # print(x, y)
            bbox_idx = get_xy_bbox_idx(x, y)
            if bbox_idx is not None:
                if bbox_idx not in select_image_bbox_idxes:
                    select_image_bbox_idxes.add(bbox_idx)
                else:
                    select_image_bbox_idxes.remove(bbox_idx)


def display_video(video_path, frame_num=None, win_name='video'):
    cap = cv2.VideoCapture(video_path)
    is_pause = False
    frame_id = 0
    while True:
        grabbed, image = cap.read()

        if not grabbed:
            break

        cv2.putText(image, 'Frame: {}'.format(frame_id+1), (32, 32), cv2.FONT_HERSHEY_COMPLEX, 1.0, (255, 0, 0))

        if (image.shape[0] != DISPLAY_HEIGHT) or (image.shape[1] != DISPLAY_WIDTH):
            image = cv2.resize(image, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
        cv2.imshow(win_name, image)
        if ((frame_num is not None) and (frame_id == frame_num)) or is_pause:
            key = cv2.waitKey()
        else:
            key = cv2.waitKey(40)
        if key == ord('g') or key == ord('G'):
            is_pause = False
        elif key == ord('z') or key == ord('Z'):
            is_pause = True
        elif key == ord('o') or key == ord('O'):
            break

        frame_id += 1

    cap.release()
    cv2.destroyWindow(win_name)
    print(video_path)


def play_video(video_name, frame_num=None, win_name='video'):
    for video_base_path in video_base_paths:
        video_path = os.path.join(video_base_path, video_name)
        if not os.path.exists(video_path):
            continue
        display_video(video_path, frame_num=frame_num, win_name=win_name)


def image_resize(img, img_width, img_height, mode=1):
    if mode == 0:
        img = cv2.resize(img, (img_width, img_height))
    elif mode == 1:
        w, h = img.shape[1], img.shape[0]
        wscale = img_width / w
        hscale = img_height / h
        scale = min(wscale, hscale)
        w = int(round(w * scale))
        h = int(round(h * scale))
        image = cv2.resize(img, (w, h))
        img = np.full((img_height, img_width, image.shape[2]), fill_value=0, dtype=np.uint8)
        x = (img_width - w) // 2
        y = (img_height - h) // 2
        img[y:y + h, x:x + w, :] = image
    return img


def show_images(image_dir, file_names, dis_text='', win_name='same_image', wait=False, h_image_count=10, v_image_count=10):
    display_image_count = h_image_count * v_image_count
    files = []
    total_image_count = len(file_names)
    last_idx = total_image_count - 1
    # if total_image_count > 0:
    #     print('{}: {}'.format(win_name, file_names[0]))
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            if file_name is None:
                img = np.zeros((img_height, img_width, 3), dtype=np.uint8)
            else:
                file_path = os.path.join(image_dir, file_name.strip())
                img = cv2.imread(file_path)
                img = image_resize(img, img_width, img_height)
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (2, 12),
                    cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return
        else:
            break


def read_file_list(file_list, label_file_path=None):
    file_labels = {}
    if label_file_path is not None:
        if os.path.exists(label_file_path):
            file_count = 0
            with open(label_file_path, 'r') as file:
                for line in file.readlines():
                    file_count += 1
                    if file_count % 10000 == 0:
                        print('Read Label File Count: {}'.format(file_count))
                    lines = line.strip().split()
                    file_path = lines[0]
                    label = int(lines[1])
                    file_labels[file_path] = label
            print('Read Label File Count: {}'.format(file_count))

    faceid_file_selects = OrderedDict()
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_154308/00010_000.jpg 1737
            line = line.strip().split()
            file_path = line[0]
            label = file_labels.get(file_path, None)
            if label is None:
                if len(line) > 1:
                    label = int(line[1])
                else:
                    label = 0

            # faceid = os.path.dirname(file_path)
            # file_name = os.path.basename(file_path)
            faceid = file_path.split('/')[-4]
            file_name = file_path
            select = 0

            file_selects = faceid_file_selects.get(faceid, None)
            if file_selects is None:
                faceid_file_selects[faceid] = [[file_name], [select], [label]]
            else:
                file_selects[0].append(file_name)
                file_selects[1].append(select)
                file_selects[2].append(label)

        print('Read File Count: {}'.format(file_count))
    faceid_list = list(faceid_file_selects.keys())
    # faceid_list.sort()
    return faceid_list, faceid_file_selects


def save_process(select_file_list_map, cur_id_idx_file_path, cur_id_idx, label_file_path=None, label_change=False):
    if (label_file_path is not None) and (label_file_path != '') and label_change:
        with open(label_file_path, 'w') as file:
            for select_id, file_selects in select_file_list_map.items():
                select_file_names, labels = file_selects[0], file_selects[2]
                for select_file_name, label in zip(select_file_names, labels):
                    label_info = '{} {}\n'.format(select_file_name, label)
                    # print(label_info)
                    file.write(label_info)
        os.system('chmod a+wr {}'.format(label_file_path))

    with open(cur_id_idx_file_path, 'w') as file:
        process_info = '{}\n'.format(cur_id_idx)
        file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_id_idx_file_path))


def filter_select_face(output_dir, file_list, filter_list):
    filter_file_set = set()
    with open(filter_list, 'r') as file:
        for line in file.readlines():
            filter_file_set.add(line)

    out_file_path = os.path.join(output_dir, 'filter_' + os.path.basename(file_list))
    out_file = open(out_file_path, 'w')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Write File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            if line in filter_file_set:
                continue

            out_file.write(line)
        print('Write File Count: {}'.format(file_count))
    out_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))


def read_same_image(same_image_list):
    video_images_dict = {}

    with open(same_image_list, 'r') as file:
        facid_images = {}
        video_faceids = {}
        for line in file.readlines():
            # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_214720/00026_000.jpg 5
            lines = line.strip().split()
            image_name = lines[0]
            faceid = int(lines[1])
            video_name = os.path.dirname(image_name)
            images = facid_images.get(faceid, None)
            if images is None:
                facid_images[faceid] = [image_name]
            else:
                images.append(image_name)
            faceids = video_faceids.get(video_name, None)
            if faceids is None:
                video_faceids[video_name] = {faceid}
            else:
                faceids.add(faceid)

        for video_name, faceids in video_faceids.items():
            images = []
            for faceid in faceids:
                images.extend(facid_images[faceid])
            images.sort()
            video_images_dict[video_name] = images
    return video_images_dict


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes, start_modify_bbox, start_mouse_xy, end_mouse_xy, \
        start_modify_bbox2, start_mouse_xy2, end_mouse_xy2, current_mouse_xy
    parser = argparse.ArgumentParser(description="Select Image Merge")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/yery/Kaso/data/ErisedPersonGender", help="path to image file", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/yery/Kaso/data/ErisedPersonGender/typical_file_faceid_list.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/yery/Kaso/data/ErisedPersonGender", help="path to mot reid image info", type=str
    )
    parser.add_argument(
        "--filter_list", default="/rootfs/media/yery/Kaso/data/ErisedPersonGender/select_image_label_merge.txt", help="path to filter file list", type=str
    )
    parser.add_argument(
        "--same_image_list", help="path to same image list", type=str
    )
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=5*60, help='save time, second ')
    parser.add_argument("--filter", action='store_true', help="filter select face")
    args = parser.parse_args()

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    if args.filter:
        filter_select_face(output_dir, args.file_list, args.filter_list)
        return

    video_images_dict = {}
    if args.same_image_list is not None:
        video_images_dict = read_same_image(args.same_image_list)

    display_image_count = h_image_count * v_image_count

    label_file_path = os.path.join(output_dir, 'select_image_label_merge.txt')
    id_list, select_file_list_map = read_file_list(args.file_list, label_file_path)
    backup_select_file_list_map = copy.deepcopy(select_file_list_map)

    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    id_count = len(id_list)
    last_id = id_list[-1]
    last_id_idx = id_count - 1

    cur_id_idx = 0
    cur_id_idx_file_path = os.path.join(output_dir, 'cur_id_idx.txt')
    if os.path.exists(cur_id_idx_file_path):
        with open(cur_id_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_id_idx = int(line.strip())
                break

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectImageMerge')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_label_file_path = os.path.join(tmp_output_dir, 'select_image_label_merge.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, 'cur_id_idx.txt')

    select_change = False
    label_change = False
    start_save_time = time.time()
    while cur_id_idx <= last_id_idx:
        try:
            cur_trk_id = id_list[cur_id_idx]

            print('{} {}'.format(cur_trk_id, last_id))
            file_list, select_list, label_list = select_file_list_map[cur_trk_id]
            _, _, backup_label_list = backup_select_file_list_map[cur_trk_id]
            last_idx = len(file_list) - 1
            start_idx = 0
            files = []
            # select_list = [0]*len(label_list)

            for idx, file_name in enumerate(file_list):
                files.append(file_name)
                file_count = len(files)
                if file_count == 1:
                    start_idx = idx
                if file_count < display_image_count and (idx < last_idx):
                    continue
                imgs = []
                for file_name in files:
                    # file_path = os.path.join(image_dir, cur_trk_id, file_name)
                    file_path = os.path.join(image_dir, file_name)
                    img = cv2.imread(file_path)
                    img = image_resize(img, img_width, img_height)
                    imgs.append(img)

                image_count = len(imgs)
                images = []
                x_off = 0
                y_off = 0
                image_bboxes.clear()
                select_image_bbox_idxes.clear()
                for select_idx in range(0, image_count):
                    if 0 != select_list[start_idx+select_idx]:
                        select_image_bbox_idxes.add(select_idx)
                for j, img in enumerate(imgs):
                    if j % h_image_count == 0:
                        x_off = 0
                        if j != 0:
                            y_off += img_height
                        h_images = [img]
                    else:
                        x_off += img_width
                        h_images.append(img)
                    image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                    if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                        images.append(np.concatenate(h_images, axis=1))

                if len(images) > 1:
                    w0 = images[0].shape[1]
                    h0 = images[0].shape[0]
                    w1 = images[-1].shape[1]
                    h1 = images[-1].shape[0]
                    if (w0 != w1) or (h0 != h1):
                        # image = np.empty_like(images[0])
                        image = np.zeros_like(images[0])
                        image[0:h1, 0:w1, :] = images[-1]
                        images[-1] = image
                draw_image = np.vstack(images)

                change_trk_id = False
                is_exit = False
                label_str = ''
                while True:
                    image = draw_image.copy()
                    if start_modify_bbox:
                        draw_rect(image, start_mouse_xy, end_mouse_xy)
                    if start_modify_bbox2:
                        draw_rect(image, start_mouse_xy2, end_mouse_xy2)

                    for bbox_idx in range(0, image_count):
                        select_idx = start_idx + bbox_idx
                        if bbox_idx in select_image_bbox_idxes:
                            select_list[select_idx] = 1
                        else:
                            select_list[select_idx] = 0
                    select_label_set = set()
                    for start_select_idx, is_select in enumerate(select_list):
                        if not is_select:
                            continue
                        label = label_list[start_select_idx]
                        select_label_set.add(label)

                    for bbox_idx in range(0, image_count):
                        bbox = image_bboxes[bbox_idx]
                        label = label_list[start_idx + bbox_idx]
                        if label in select_label_set:
                            cv2.circle(image, (bbox[0] + 14, bbox[1] + 14), 8, (0, 0, 255), -1)
                        cv2.putText(image, '%d' % (label, ), (bbox[0] + 4, bbox[3] - 4), cv2.FONT_HERSHEY_PLAIN, 1.5, (0, 0, 255), thickness=2)
                    mouse_xy = current_mouse_xy
                    if mouse_xy is None:
                        cur_mouse_bbox_idx = None
                    else:
                        cur_mouse_bbox_idx = get_xy_bbox_idx(*mouse_xy)
                    if cur_mouse_bbox_idx is not None:
                        cur_label = label_list[start_idx + cur_mouse_bbox_idx]
                        for bbox_idx in range(0, image_count):
                            label = label_list[start_idx + bbox_idx]
                            if label != cur_label:
                                continue
                            bbox = image_bboxes[bbox_idx]
                            cv2.circle(image, (bbox[0] + 14, bbox[1] + 32), 8, (255, 0, 0), -1)
                    cv2.putText(image, '%d/%d %d/%d' % (cur_id_idx+1, last_id_idx+1, idx + 1, last_idx + 1), (8, 24), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 255, 0), thickness=2)
                    cv2.imshow('image', image)
                    key = cv2.waitKey(50)
                    if (key == 27) or (key == ord('q')) or (key == ord('Q')):
                        is_exit = True
                        break
                    elif key == ord(' '):  # 同一ID的下一页
                        select_change = True
                        break
                    elif key == ord('a') or key == ord('A'):  # 上一个ID
                        for start_select_idx, _ in enumerate(select_list):
                            select_list[start_select_idx] = 0
                        select_change = True
                        if cur_id_idx > 0:
                            cur_id_idx -= 1
                            change_trk_id = True
                        break
                    elif key == ord('d') or key == ord('D'):  # 下一个ID
                        for start_select_idx, _ in enumerate(select_list):
                            select_list[start_select_idx] = 0
                        select_change = True
                        change_trk_id = True
                        cur_id_idx += 1
                        break
                    elif key == ord('e') or key == ord('E'):  # Merge选中图片的Label
                        select_label_set = set()
                        for start_select_idx, is_select in enumerate(select_list):
                            select_list[start_select_idx] = 0
                            if not is_select:
                                continue
                            label = label_list[start_select_idx]
                            select_label_set.add(label)
                        new_label = None
                        for start_select_idx, label in enumerate(label_list):
                            if label not in select_label_set:
                                continue
                            if new_label is None:
                                new_label = label
                            label_list[start_select_idx] = new_label
                        select_image_bbox_idxes.clear()
                        label_change = True
                        continue
                    elif key == ord('w') or key == ord('W'):  # 全不选
                        for start_select_idx, _ in enumerate(select_list):
                            select_list[start_select_idx] = 0
                        select_image_bbox_idxes.clear()
                        continue
                    elif key == ord('r') or key == ord('R'):  # 恢复当前鼠标所在图片的Label
                        if cur_mouse_bbox_idx is not None:
                            label_list[start_idx + cur_mouse_bbox_idx] = backup_label_list[start_idx + cur_mouse_bbox_idx]
                            label_change = True
                        continue
                    elif key == ord('f') or key == ord('F'):  # 播放视频
                        if cur_mouse_bbox_idx is not None:
                            image_file_path = file_list[start_idx + cur_mouse_bbox_idx]
                            print('F:', image_file_path)
                            video_name = os.path.dirname(image_file_path) + '.avi'
                            image_name = os.path.basename(image_file_path)
                            frame_num = int(image_name.split('_')[0])
                            play_video(video_name, frame_num, win_name='P{}'.format(cur_mouse_bbox_idx))
                        continue
                    elif key == ord('v') or key == ord('V'):  # 当前鼠标所在图片
                        if cur_mouse_bbox_idx is not None:
                            _file_name = file_list[start_idx + cur_mouse_bbox_idx]
                            # image_file_path = os.path.join(image_dir, cur_trk_id, _file_name)
                            image_file_path = os.path.join(image_dir, _file_name)
                            print('V:', image_file_path)
                            _image = cv2.imread(image_file_path)
                            cv2.imshow('{}'.format(cur_mouse_bbox_idx), _image)
                        continue
                    elif key == ord('k') or key == ord('K'):  # 当前鼠标所在图片
                        if cur_mouse_bbox_idx is not None:
                            _file_name = file_list[start_idx + cur_mouse_bbox_idx]
                            image_names = video_images_dict.get(os.path.dirname(_file_name), None)
                            if image_names is None:
                                continue
                            show_images(image_dir, image_names, win_name='image{}'.format(start_idx + cur_mouse_bbox_idx))
                        continue
                    elif (ord('0') <= key <= ord('9')) or (key == ord('-')):  # 框输入新Label
                        if label_str == '':
                            label_str = chr(key)
                        else:
                            label_str += chr(key)
                    elif (key == ord('l')) or (key == ord('L')):  # 修改这个ID所有图片的Label
                        if label_str != '':
                            label = int(label_str)
                            _count = len(label_list)
                            for _idx in range(_count):
                                label_list[_idx] = label
                            label_change = True
                            label_str = ''
                    elif (key == ord('c')) or (key == ord('C')):  # 修改这当前鼠标所在图片的Label
                        if label_str != '':
                            if cur_mouse_bbox_idx is not None:
                                label = int(label_str)
                                label_list[start_idx + cur_mouse_bbox_idx] = label
                                label_change = True
                            label_str = ''
                    elif key == ord('s') or key == ord('S'):  # Save
                        if select_change or label_change:
                            save_process(select_file_list_map, cur_id_idx_file_path, cur_id_idx, label_file_path, label_change)
                            if args.use_temp:
                                save_process(select_file_list_map, tmp_cur_id_idx_file_path, cur_id_idx, tmp_label_file_path, label_change)
                            label_change = True
                        start_save_time = time.time()
                        select_change = False
                        continue
                    if args.use_temp:
                        # 定时保存
                        if time.time() - start_save_time >= save_time:
                            if select_change or label_change:
                                save_process(select_file_list_map, tmp_cur_id_idx_file_path, cur_id_idx, tmp_label_file_path, label_change)
                                select_change = False
                            start_save_time = time.time()
                files.clear()
                if is_exit:
                    save_process(select_file_list_map, cur_id_idx_file_path, cur_id_idx, label_file_path, label_change)
                    return
                if change_trk_id:
                    break
        except Exception as e:
            print(e)
    save_process(select_file_list_map, cur_id_idx_file_path, cur_id_idx, label_file_path, label_change)
    print('finish!')


if __name__ == '__main__':
    main()
